#!/usr/bin/env python3
"""
A script to generate quantum circuit diagrams using ProjectQ.
This script provides functions to create and export quantum circuits to LaTeX format.
"""

from projectq import MainEngine
from projectq.backends import CircuitDrawer
from projectq.ops import H, CNOT, X, Y, Z, Measure, All, BasicGate
import numpy as np
import os

class UGate(BasicGate):
    """Custom unitary gate for drawing."""
    def __init__(self):
        BasicGate.__init__(self)
        self.name = "U"

    def __str__(self):
        return "U"

class XGate(BasicGate):
    """Custom X gate for drawing."""
    def __init__(self):
        BasicGate.__init__(self)
        self.name = "X"

    def __str__(self):
        return "X"

def create_and_draw_circuit(num_qubits=2, gates=None):
    """
    Create a quantum circuit and export it to LaTeX format.
    
    Args:
        num_qubits (int): Number of qubits in the circuit
        gates (list): List of tuples (gate_type, qubit, position)
        
    Returns:
        str: LaTeX code for the circuit
    """
    # Create a new circuit drawer
    drawing_engine = CircuitDrawer()
    
    # Create a new main engine
    eng = MainEngine(drawing_engine)
    
    # Allocate qubits
    qubits = eng.allocate_qureg(num_qubits)
    
    # Create gate instances
    U = UGate()
    X_custom = XGate()
    
    # Apply gates if specified
    if gates is None:
        # Default circuit: Hadamard on first qubit, CNOT between first and second qubit
        H | qubits[0]
        CNOT | (qubits[0], qubits[1])
    else:
        for gate_type, qubit, _ in gates:
            if gate_type == "H":
                H | qubits[qubit]
            elif gate_type == "CNOT":
                CNOT | (qubits[qubit], qubits[qubit + 1])
            elif gate_type == "X":
                X_custom | qubits[qubit]  # Use custom X gate
            elif gate_type == "Y":
                Y | qubits[qubit]
            elif gate_type == "Z":
                Z | qubits[qubit]
            elif gate_type == "U":
                U | qubits[qubit]
    
    # Measure all qubits
    All(Measure) | qubits
    
    # Flush the engine
    eng.flush()
    
    # Get the LaTeX code
    latex_code = drawing_engine.get_latex()
    
    # Deallocate qubits
    for qubit in qubits:
        del qubit
    
    return latex_code

def save_to_latex(latex_code, filename="circuit.tex"):
    """
    Save the LaTeX code to a file.
    
    Args:
        latex_code (str): The LaTeX code to save
        filename (str): The name of the file to save to
    """
    with open(filename, "w") as f:
        f.write(latex_code)

def create_evolution_circuit(num_qubits=2):
    """
    Create a circuit demonstrating quantum evolution with unitary operator U.
    
    Args:
        num_qubits (int): Number of qubits in the circuit
        
    Returns:
        str: LaTeX code for the circuit
    """
    drawing_engine = CircuitDrawer()
    eng = MainEngine(drawing_engine)
    
    # Allocate qubits
    qubits = eng.allocate_qureg(num_qubits)
    
    # Create gate instances
    U = UGate()
    X_custom = XGate()
    
    # Apply evolution operator U
    U | qubits[0]
    
    # Apply CNOT for entanglement
    CNOT | (qubits[0], qubits[1])
    CNOT | (qubits[1], qubits[2])
    

    # Measure first qubit
    Measure | qubits[0]
    
    # # Apply X gate
    # X_custom | qubits[1]
    
    # Measure second qubit
    Measure | qubits[1]
    Measure | qubits[2]
    
    eng.flush()
    latex_code = drawing_engine.get_latex()
    
    for qubit in qubits:
        del qubit
    
    return latex_code

def create_time_evolution_circuit(num_qubits=2, time_steps=3):
    """
    Create a circuit showing time evolution with multiple unitary operators.
    
    Args:
        num_qubits (int): Number of qubits in the circuit
        time_steps (int): Number of time evolution steps
        
    Returns:
        str: LaTeX code for the circuit
    """
    drawing_engine = CircuitDrawer()
    eng = MainEngine(drawing_engine)
    
    # Allocate qubits
    qubits = eng.allocate_qureg(num_qubits)
    
    # Create gate instances
    U = UGate()
    X_custom = XGate()
    
    # Apply time evolution
    for _ in range(time_steps):
        U | qubits[0]
        CNOT | (qubits[0], qubits[1])
        X_custom | qubits[1]
    
    # Measure
    All(Measure) | qubits
    
    eng.flush()
    latex_code = drawing_engine.get_latex()
    
    for qubit in qubits:
        del qubit
    
    return latex_code

if __name__ == "__main__":
    # # Example 1: Basic circuit with Hadamard and CNOT
    # print("Generating basic circuit...")
    # circuit = create_and_draw_circuit(2)
    # save_to_latex(circuit, "basic_circuit.tex")
    
    # Example 2: Evolution circuit
    print("Generating evolution circuit...")
    circuit = create_evolution_circuit(3)
    save_to_latex(circuit, "result.tex")
    
    # # Example 3: Time evolution circuit
    # print("Generating time evolution circuit...")
    # circuit = create_time_evolution_circuit(2, time_steps=3)
    # save_to_latex(circuit, "time_evolution_circuit.tex")
    
    # print("LaTeX code has been generated and saved to basic_circuit.tex, evolution_circuit.tex, and time_evolution_circuit.tex") 